# -*- coding: utf-8 -*-
# ===========================================
# @Time    : 2021/9/13 11:30 
# @Author  : shutao
# @FileName: base_net.py
# @remark  : 
# 
# @Software: PyCharm
# Github 　： https://github.com/NameLacker
# ===========================================

import numpy as np
import paddle
import paddle.nn as nn

import paddle.nn.functional as F


def _get_feature_map(maps, feature_map_name):
    """
    Extract feature map.
    Args:
        feature_map_name: Feature map name
        maps: The output feature space vector of a certain layer of the network.
            x.shape = [N, C, H, W]

    Returns:
        x: Return to displayable feature map.
            feature_map: Displayable feature map
            feature_name: Feature map name
    """
    N, C, H, W = maps.shape
    out_maps = paddle.zeros((N, H, W), dtype=paddle.int32)
    for idx, map in enumerate(maps):
        map = paddle.sum(map, 0).squeeze()
        t_max = paddle.max(map)
        t_min = paddle.min(map)
        area = t_max - t_min
        map = (map - t_min) * 255 / area
        out_maps[idx] = paddle.cast(map, paddle.int32)

    return {"feature_map": out_maps,
            "feature_name": feature_map_name}


class BaseNet(nn.Layer):
    """Base class of any neural network"""

    def __init__(self):
        super(BaseNet, self).__init__()

        self.feature_maps = []  # Used to store feature maps.

        self.sigmoid = nn.Sigmoid()
        self.flatten = nn.Flatten()
        self.relu = nn.ReLU()

    def _get_feature_maps_name(self):
        map_names = []
        for f in self.feature_maps:
            map_names.append(f["feature_name"])
        return map_names

    def append_maps(self, x: paddle.Tensor, feature_map_name: str) -> None:
        """
        Add feature map
        Args:
            feature_map_name: Feature map name
            x: Network layer output feature map

        Returns:
            None
        """
        if self.training:
            feature = _get_feature_map(x, feature_map_name)
            if len(self.feature_maps) == 0:
                self.feature_maps.append(feature)
            elif feature["feature_name"] not in self._get_feature_maps_name():
                self.feature_maps.append(feature)
            else:
                for idx, f in enumerate(self.feature_maps):
                    if f["feature_name"] == feature["feature_name"]:
                        self.feature_maps[idx] = feature

    def get_feature_maps(self):
        """
        Return to displayable network feature map
        Returns:
            feature map
        """
        if len(self.feature_maps):
            return self.feature_maps
        else:
            return None

    def forward(self, *inputs, **kwargs):
        raise NotImplementedError
